#coding=utf-8

import moviepy.editor as mp
from moviepy.editor import VideoFileClip
from PIL import Image
import os
from moviepy.editor import *
import cv2
import sys
import random
import numpy as np
import matplotlib.pyplot as plt
cmd = None
debug = False
seed = 20160930
oldseed = False
alpha = 5.0
water_mask_path = 'wm.png'
temp_cover = 'temp_mask'

# 视频增加水印
def set_mask(cover):
    video = mp.VideoFileClip("./old.mp4")

    logo = (mp.ImageClip(cover)
            .set_start(10)
            .set_duration(2)
            .resize(video.size)
            )  # 水印的位置

    final = mp.CompositeVideoClip([video, logo])
    # mp4文件默认用libx264编码， 比特率单位bps
    final.write_videofile("./new.mp4")  

# 加密水印 fn1 原图，fn2 水印图,fn3 输出
def encode_mask(fn1,fn2,fn3):
    img = cv2.imread(fn1)
    wm = cv2.imread(fn2)

    # print img.shape # 高, 宽, 通道
    h, w = img.shape[0], img.shape[1]
    hwm = np.zeros((int(h * 0.5), w, img.shape[2]))
    assert hwm.shape[0] > wm.shape[0]
    assert hwm.shape[1] > wm.shape[1]
    hwm2 = np.copy(hwm)
    for i in range(wm.shape[0]):
        for j in range(wm.shape[1]):
            hwm2[i][j] = wm[i][j]

    if oldseed: random.seed(seed,version=1)
    else: random.seed(seed)
    m, n = list(range(hwm.shape[0])), list(range(hwm.shape[1]))
    if oldseed:
        random.shuffle(m,random=random.random)
        random.shuffle(n,random=random.random)
    else:
        random.shuffle(m)
        random.shuffle(n)

    for i in range(hwm.shape[0]):
        for j in range(hwm.shape[1]):
            hwm[i][j] = hwm2[m[i]][n[j]]

    rwm = np.zeros(img.shape)
    for i in range(hwm.shape[0]):
        for j in range(hwm.shape[1]):
            rwm[i][j] = hwm[i][j]
            rwm[rwm.shape[0] - i - 1][rwm.shape[1] - j - 1] = hwm[i][j]

   

    f1 = np.fft.fft2(img)
    f2 = f1 + alpha * rwm
    _img = np.fft.ifft2(f2)

    

    img_wm = np.real(_img)
    # img_wm = bgr_to_rgba(img_wm)
    # print (img_wm)
    assert cv2.imwrite(fn3, img_wm, [int(cv2.IMWRITE_PNG_COMPRESSION), 3])

# 解码
def decode_mask(fn1,fn2,fn3):
    img = cv2.imread(fn1)
    img_wm = cv2.imread(fn2)

    if oldseed: random.seed(seed,version=1)
    else: random.seed(seed)
    m, n = list(range(int(img.shape[0] * 0.5))), list(range(img.shape[1]))
    if oldseed:
        random.shuffle(m,random=random.random)
        random.shuffle(n,random=random.random)
    else:
        random.shuffle(m)
        random.shuffle(n)

    f1 = np.fft.fft2(img)
    f2 = np.fft.fft2(img_wm)

    rwm = (f2 - f1) / alpha
    rwm = np.real(rwm)

    
    wm = np.zeros(rwm.shape)
    for i in range(int(rwm.shape[0] * 0.5)):
        for j in range(rwm.shape[1]):
            wm[m[i]][n[j]] = np.uint8(rwm[i][j])
    for i in range(int(rwm.shape[0] * 0.5)):
        for j in range(rwm.shape[1]):
            wm[rwm.shape[0] - i - 1][rwm.shape[1] - j - 1] = wm[i][j]
    assert cv2.imwrite(fn3, wm)



def video_2_imgs(fg_video_path,output):
    clip = VideoFileClip(fg_video_path).subclip(10,10.1)
    i = 0
    if not os.path.exists(output):
        os.mkdir(output)
    for frame in clip.iter_frames():

        im = Image.fromarray(frame)
        # im.resize(size)
        # im = im.convert('RGBA')  
        print (im)
        im.save(output+"/%07d.png" % i)
        i = i+1

if __name__ == '__main__':
    # 加密
    print ('Start encode')
    # 获得原视频的抽帧
    video_2_imgs('old.mp4','old_frame_image')
    # 给原视频抽帧的图加水印
    if not os.path.exists(temp_cover):
        os.mkdir(temp_cover)
    encode_mask('./old_frame_image/0000000.png','wm.png', temp_cover+'/cover.png')
    # 将加过水印的图覆盖到原视频
    set_mask(temp_cover+'/cover.png')
    # 将加过水印的图抽帧
    if os.path.exists('new.mp4'):
        video_2_imgs('new.mp4','new_frame_image')
    
    # 解码
    decode_mask("./old_frame_image/0000000.png","./new_frame_image/0000000.png",'decode.png')
    print ('Blind Mask Done')



